"""
Code for IceFormer: Accelerated Inference with Long-Sequence Transformers on CPUs

This code implements the method described in the IceFormer paper,
which can be found at https://openreview.net/forum?id=6RR3wU4mSZ

This file is a part of the Dynamic Continuous Indexing reference
implementation.

This Source Code Form is subject to the terms of the Mozilla Public
License, v. 2.0. If a copy of the MPL was not distributed with this
file, You can obtain one at https://mozilla.org/MPL/2.0/.

Copyright (C) 2024    Yuzhen Mao, Ke Li
"""

import os
import sys
from setuptools import setup, Extension
from setuptools.command.build_ext import build_ext


# Helper function to prompt for multithreading support
def enable_multithreading():
    answer = input("Enable Multithreading? (Y/N)\n")
    return answer.lower().startswith("y")


# Gather the source and header files
dci_sources = [
    'src/dci.c', 'src/py_dci.c', 'src/util.c',
    'src/hashtable_i.c', 'src/hashtable_d.c',
    'src/btree_i.c', 'src/btree_p.c',
    'src/hashtable_p.c', 'src/hashtable_pp.c',
]
dci_headers = [
    'include/dci.h', 'include/util.h', 'include/hashtable_i.h',
    'include/hashtable_d.h', 'include/btree_i.h', 'include/btree_p.h',
    'include/hashtable_p.h', 'include/hashtable_pp.h',
]


# Build the extension module
def build_ext_modules():
    from numpy import get_include
    include_dirs = [get_include(), 'include']

    # Base compiler arguments
    extra_compile_args = []
    extra_link_args = []

    # Add OpenMP support if requested
    if enable_multithreading():
        # 移除不必要的 -arch 参数
        extra_compile_args = ['-mcpu=native']
        # extra_compile_args = ['-fopenmp', '-DUSE_OPENMP', '-mcpu=native']
        extra_link_args.append('-lgomp')

    return [
        Extension(
            '_dci',
            sources=dci_sources,
            include_dirs=include_dirs,
            extra_compile_args=extra_compile_args,
            extra_link_args=extra_link_args,
        )
    ]


# Read long description
long_description = """
Dynamic Continuous Indexing (DCI) is a family of randomized algorithms for
exact k-nearest neighbour search that overcomes the curse of dimensionality.
Its query time complexity is linear in ambient dimensionality and sublinear
in intrinsic dimensionality. ``dciknn`` is a python package that contains
the reference implementation of a modified version of DCI 
and a convenient Python interface which can be used to accelerate Transformers. 

``dciknn`` requires ``NumPy``. 
"""

# Setup configuration
setup(
    name="dciknn",
    version="0.1.0",
    description="Dynamic Continuous Indexing reference implementation.",
    author="Yuzhen Mao, Ke Li",
    author_email="yuzhenm@sfu.ca",
    url="https://yuzhenmao.github.io/IceFormer/",
    license="Mozilla Public License 2.0",
    classifiers=[
        'Development Status :: 4 - Beta',
        'Environment :: Console',
        'Operating System :: OS Independent',
        'Intended Audience :: Developers',
        'Intended Audience :: Science/Research',
        'License :: OSI Approved :: Mozilla Public License 2.0 (MPL 2.0)',
        'Programming Language :: Python :: 3.6',
        'Programming Language :: Python :: 3.7',
        'Programming Language :: Python :: 3.8',
        'Programming Language :: Python :: 3.9',
        'Programming Language :: Python :: 3.10',
        'Programming Language :: Python :: 3.11',
        'Topic :: Scientific/Engineering :: Mathematics',
        'Topic :: Software Development :: Libraries :: Python Modules',
    ],
    install_requires=['numpy'],
    long_description=long_description,
    ext_modules=build_ext_modules(),
    packages=["dciknn"],
)